import json
import time
import base64
from functools import wraps
from Crypto.Cipher import AES
from Crypto.Random import get_random_bytes
from Crypto.Util.Padding import pad, unpad
from fastapi import Request, HTTPException

class EncryptionManager:
    def __init__(self, secret_key: str = "qwertyuiop1234567654zaqwsxcderfv"):
        self.secret_key = secret_key.encode('utf-8')[:32].ljust(32, b'\0')
    
    def aes_encrypt(self, data: str) -> str:
        """AES256加密"""
        iv = get_random_bytes(16)
        cipher = AES.new(self.secret_key, AES.MODE_CBC, iv)
        padded_data = pad(data.encode('utf-8'), AES.block_size)
        encrypted = cipher.encrypt(padded_data)
        return base64.b64encode(iv + encrypted).decode('utf-8')
    
    def aes_decrypt(self, encrypted_data: str) -> str:
        """AES256解密"""
        try:
            encrypted_bytes = base64.b64decode(encrypted_data)
            iv = encrypted_bytes[:16]
            encrypted = encrypted_bytes[16:]
            cipher = AES.new(self.secret_key, AES.MODE_CBC, iv)
            decrypted = unpad(cipher.decrypt(encrypted), AES.block_size)
            return decrypted.decode('utf-8')
        except Exception as e:
            raise ValueError(f"解密失败: {str(e)}")

# 全局加密管理器实例
encryption_manager = EncryptionManager()

def require_encryption(required_fields=None, optional_fields=None, skip_encryption=False):
    def decorator(func):
        @wraps(func)
        async def wrapper(request: Request, *args, **kwargs):
            if skip_encryption:
                return await func(request, *args, **kwargs)
            
            try:
                # 获取请求数据
                if request.method == "GET":
                    encrypted_data = request.query_params.get("encrypted_data")
                else:
                    body = await request.body()
                    print(f"[DEBUG] 原始请求体: {body.decode('utf-8')}")
                    if body:
                        request_data = json.loads(body)
                        encrypted_data = request_data.get("encrypted_data")
                        print(f"[DEBUG] 提取的加密数据: {encrypted_data[:50]}...")
                    else:
                        encrypted_data = None
                
                if not encrypted_data:
                    raise HTTPException(status_code=400, detail="缺少加密数据")
                
                # 解密数据
                decrypted_json = encryption_manager.aes_decrypt(encrypted_data)
                decrypted_data = json.loads(decrypted_json)
                print(f"[DEBUG] 第一次解密后数据: {decrypted_data}")
                
                # 检查是否存在双重加密（如果解密后仍然包含encrypted_data字段）
                if 'encrypted_data' in decrypted_data and len(decrypted_data) == 1:
                    print("[DEBUG] 检测到双重加密，进行第二次解密")
                    second_encrypted_data = decrypted_data['encrypted_data']
                    decrypted_json = encryption_manager.aes_decrypt(second_encrypted_data)
                    decrypted_data = json.loads(decrypted_json)
                    print(f"[DEBUG] 第二次解密后数据: {decrypted_data}")
                
                # 验证时间戳（但不要求timestamp必须存在）
                verification_timestamp = decrypted_data.get("timestamp")
                if verification_timestamp:
                    current_time = int(time.time())
                    time_diff = abs(current_time - verification_timestamp)
                    print(f"[DEBUG] 时间差: {time_diff}秒")
                    if time_diff > 300:
                        raise HTTPException(status_code=400, detail="请求已过期")
                
                # 验证必需字段
                if required_fields:
                    for field in required_fields:
                        if field not in decrypted_data:
                            print(f"[DEBUG] 缺少字段: {field}, 可用字段: {list(decrypted_data.keys())}")
                            raise HTTPException(status_code=400, detail=f"缺少必需字段: {field}")
                
                # 创建最终数据副本，始终移除timestamp（因为它只用于验证）
                final_data = decrypted_data.copy()
                if "timestamp" in final_data:
                    del final_data["timestamp"]
                    print(f"[DEBUG] 已移除验证用的timestamp，最终数据: {final_data}")
                
                request.state.decrypted_data = final_data
                return await func(request, *args, **kwargs)
                
            except json.JSONDecodeError as e:
                print(f"[DEBUG] JSON解析错误: {e}")
                raise HTTPException(status_code=400, detail="无效的JSON数据")
            except ValueError as e:
                print(f"[DEBUG] 解密错误: {e}")
                raise HTTPException(status_code=400, detail=str(e))
            except Exception as e:
                print(f"[DEBUG] 其他错误: {e}")
                raise HTTPException(status_code=500, detail=f"服务器错误: {str(e)}")
        
        return wrapper
    return decorator

def get_decrypted_data(request: Request) -> dict:
    """从request中获取解密后的数据"""
    return getattr(request.state, 'decrypted_data', {})

def create_encrypted_payload(data: dict) -> dict:
    """创建加密载荷（客户端使用）"""
    # 添加时间戳
    data["timestamp"] = int(time.time())
    
    # 加密数据
    json_data = json.dumps(data, ensure_ascii=False)
    encrypted_data = encryption_manager.aes_encrypt(json_data)
    
    return {"encrypted_data": encrypted_data}


def verify_encrypted_data(encrypted_data: str, required_fields: list = None) -> dict:
    """
    验证加密数据的完整性和有效性
    
    Args:
        encrypted_data: 加密的数据字符串
        required_fields: 必需字段列表
    
    Returns:
        dict: 解密后的数据
    
    Raises:
        ValueError: 解密失败或验证失败
    """
    try:
        # 解密数据
        decrypted_json = encryption_manager.aes_decrypt(encrypted_data)
        decrypted_data = json.loads(decrypted_json)
        
        # 验证时间戳（防重放攻击）
        timestamp = decrypted_data.get("timestamp")
        if timestamp:
            current_time = int(time.time())
            if abs(current_time - timestamp) > 300:  # 5分钟有效期
                raise ValueError("请求已过期")
        
        # 验证必需字段
        if required_fields:
            for field in required_fields:
                if field not in decrypted_data:
                    raise ValueError(f"缺少必需字段: {field}")
        
        return decrypted_data
        
    except json.JSONDecodeError:
        raise ValueError("无效的JSON数据")
    except Exception as e:
        raise ValueError(f"验证失败: {str(e)}")


def test_encryption():
    """
    测试加密解密功能
    """
    print("=== 加密解密测试 ===")
    
    # 测试数据
    test_data = {
        "userPhone": "13800138000",
        "userCode": "123456",
        "message": "这是一条测试消息"
    }
    
    try:
        # 创建加密载荷
        encrypted_payload = create_encrypted_payload(test_data.copy())
        print(f"原始数据: {test_data}")
        print(f"加密后数据: {encrypted_payload['encrypted_data']}")
        
        # 验证加密数据
        decrypted_data = verify_encrypted_data(
            encrypted_payload['encrypted_data'], 
            required_fields=["userPhone", "userCode"]
        )
        print(f"解密后数据: {decrypted_data}")
        
        # 验证数据一致性
        for key, value in test_data.items():
            if decrypted_data.get(key) == value:
                print(f"✓ {key}: 验证通过")
            else:
                print(f"✗ {key}: 验证失败")
        
        print("\n=== 测试用户提供的加密数据 ===")
        user_encrypted_data = "vI5Y+FX595yiErgmfJOq7sCRRuBVu\/3RvskhHRgGBldbwTnJvmhaNRZ1ksAKBXHSMnVqhigq1UVlFNlu3YKj300PWx3G\/pG9qkojqZOW5UpWLxEaCN5yFUYhGDT5DRj+"
        
        try:
            user_decrypted_data = verify_encrypted_data(user_encrypted_data)
            print(f"用户数据解密成功: {user_decrypted_data}")
        except ValueError as e:
            print(f"用户数据解密失败: {e}")
            
    except Exception as e:
        print(f"测试失败: {e}")


if __name__ == "__main__":
    test_encryption()